import asyncio
import aiofiles
from fastapi import APIRouter, Depends, UploadFile
from fastapi.params import Form
from fastapi.responses import StreamingResponse

from ..models.user import User
from ..schemas.file import FileMetadata, FileOperationResponse, CatalogResponse
from ..services.file import *
from ..db import get_db
from ..core import USER_STORAGE_PATH

from .auth import get_current_user

file_router = APIRouter()


# 删除文件
@file_router.post("/delete", response_model=FileOperationResponse)
def delete_file(file_name: str, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    # noinspection PyBroadException
    try:
        if delete_file_by_name(db, file_name, user):
            return FileOperationResponse(success=True, message="file deleted successfully")
        else:
            return FileOperationResponse(success=False, message="file not found")
    except Exception as e:
        return FileOperationResponse(success=False, message=str(e))


@file_router.post("/apply_for_upload")
def apply_for_upload(file_meta_data: FileMetadata, user: User = Depends(get_current_user),
                     db: Session = Depends(get_db)):
    if file_meta_data.file_size + user.disk_used > user.disk_quota:
        return FileOperationResponse(success=False, message="disk quota exceeded")
    if file_exists(db, file_meta_data.file_name, user.id):
        return FileOperationResponse(success=False, message="file already exists")
    return FileOperationResponse(success=True, message="allowed")


# 上传文件
@file_router.post("/upload", response_model=FileOperationResponse)
def upload_file(file: UploadFile, file_name: str = Form(...), user: User = Depends(get_current_user),
                db: Session = Depends(get_db)):
    try:
        create_file(db, file_name, user, file.file.read())
        return FileOperationResponse(success=True, message="file uploaded successfully")
    except Exception as e:
        return FileOperationResponse(success=False, message=str(e))


# 下载文件
@file_router.get("/download")
async def download_file(file_name: str, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    file_info = get_file_by_name(db, file_name, user.id)
    file_path = Path(USER_STORAGE_PATH, file_info.folder_path, file_name)
    if not file_path.exists():
        return FileOperationResponse(success=False, message="file not found")
    delay = 0 if user.vip else 0.03  # 每次发送后的延时
    chunk_size = 1024  # 每次读取的字节数（例如1KB）

    async def file_iterator(f_p):
        async with aiofiles.open(f_p, "rb") as file:
            while chunk := await file.read(chunk_size):
                yield chunk
                await asyncio.sleep(delay)

    headers = {
        "Content-Disposition": f"attachment; filename={file_info.file_name}",
        "Content-Length": str(file_info.file_size),
        "Access-Control-Allow-Origin": "*",  # 允许所有源访问
        "Access-Control-Allow-Methods": "*",  # 允许的HTTP方法
        "Access-Control-Allow-Headers": "*",  # 允许的头部
        "Access-Control-Allow-Credentials": "true",  # 允许携带凭证
    }
    return StreamingResponse(file_iterator(file_path), headers=headers, media_type="application/octet-stream")


# 获取文件列表
@file_router.get("/list")
def list_files(user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    # noinspection PyBroadException
    try:
        files = get_files_by_owner(db, user.id)
        return CatalogResponse(success=True, files=[e.file_name for e in files])
    except Exception as e:
        return CatalogResponse(success=False, message=str(e))
